import numpy as np
import matplotlib.pyplot as plt
from PIL import Image


class ImageProcessor:
    def __init__(self, image_path):
        """初始化图像处理器"""
        # 使用PIL读取图像，转换为numpy数组
        self.original_image = np.array(Image.open(image_path))
        # 转换为灰度图
        self.gray_image = self.rgb_to_gray(self.original_image)

    def rgb_to_gray(self, rgb_image):
        """自实现RGB转灰度图
        使用加权平均法: Gray = 0.299R + 0.587G + 0.114B
        """
        r, g, b = rgb_image[:, :, 0], rgb_image[:, :, 1], rgb_image[:, :, 2]
        gray = 0.299 * r + 0.587 * g + 0.114 * b
        return gray.astype(np.uint8)

    def convolution2d(self, image, kernel):
        """自实现2D卷积运算"""
        # 获取图像和卷积核的尺寸
        img_h, img_w = image.shape
        kernel_h, kernel_w = kernel.shape

        # 计算需要的填充量
        pad_h = kernel_h // 2
        pad_w = kernel_w // 2

        # 创建填充后的图像
        padded_image = np.zeros((img_h + 2 * pad_h, img_w + 2 * pad_w))
        padded_image[pad_h:pad_h + img_h, pad_w:pad_w + img_w] = image

        # 输出图像
        output = np.zeros_like(image, dtype=np.float32)

        # 执行卷积运算
        for i in range(img_h):
            for j in range(img_w):
                output[i, j] = np.sum(
                    padded_image[i:i + kernel_h, j:j + kernel_w] * kernel
                )

        return output

    def sobel_filter(self):
        """自实现Sobel边缘检测"""
        # 定义Sobel算子
        sobel_x = np.array([[-1, 0, 1],
                            [-2, 0, 2],
                            [-1, 0, 1]])

        sobel_y = np.array([[-1, -2, -1],
                            [0, 0, 0],
                            [1, 2, 1]])

        # 分别计算x和y方向的梯度
        grad_x = self.convolution2d(self.gray_image, sobel_x)
        grad_y = self.convolution2d(self.gray_image, sobel_y)

        # 计算梯度幅值
        gradient_magnitude = np.sqrt(grad_x ** 2 + grad_y ** 2)

        # 归一化到0-255
        gradient_magnitude = (gradient_magnitude / gradient_magnitude.max() * 255).astype(np.uint8)

        return gradient_magnitude

    def custom_filter(self, kernel):
        """使用自定义卷积核进行滤波"""
        filtered = self.convolution2d(self.gray_image, kernel)
        # 归一化到0-255
        filtered = (filtered - filtered.min()) / (filtered.max() - filtered.min()) * 255
        return filtered.astype(np.uint8)

    def gaussian_blur(self):
        """高斯模糊"""
        # 创建5x5的高斯核
        kernel_size = 5
        sigma = 1.0
        k = kernel_size // 2
        x, y = np.mgrid[-k:k + 1, -k:k + 1]
        gaussian_kernel = np.exp(-(x ** 2 + y ** 2) / (2 * sigma ** 2))
        gaussian_kernel = gaussian_kernel / gaussian_kernel.sum()

        # 对灰度图像进行高斯模糊
        blurred_gray = self.convolution2d(self.gray_image, gaussian_kernel)

        return blurred_gray.astype(np.uint8)

    def edge_detection_after_gaussian(self):
        """高斯平滑后进行边缘检测"""
        # 先进行高斯模糊
        blurred = self.gaussian_blur()

        # 保存原始灰度图像
        original_gray = self.gray_image.copy()

        # 临时将模糊后的图像设置为处理图像
        self.gray_image = blurred

        # 使用现有的边缘检测方法
        sobel_result = self.sobel_filter()

        custom_kernel = np.array([[1, 0, -1],
                                  [2, 0, -2],
                                  [1, 0, -1]])
        custom_result = self.custom_filter(custom_kernel)

        # 恢复原始灰度图像
        self.gray_image = original_gray

        return sobel_result, custom_result
    def color_histogram(self):
        """自实现颜色直方图计算"""
        histograms = np.zeros((256, 3))
        # 分别计算RGB三个通道的直方图
        for channel in range(3):
            channel_data = self.original_image[:, :, channel].ravel()
            for pixel_value in channel_data:
                histograms[pixel_value, channel] += 1

        return histograms
    def extract_texture_features(self):
        """提取纹理特征并归一化
        包括：均值、标准差、能量、熵、对比度
        所有特征都归一化到[0,1]区间
        """
        features = {}

        # 图像归一化到[0,1]
        img_normalized = self.gray_image / 255.0

        # 计算均值 (已经归一化到[0,1])
        features['mean'] = np.mean(img_normalized)

        # 计算标准差 (已经归一化到[0,1])
        features['std'] = np.std(img_normalized)

        # 计算能量并归一化
        energy = np.sum(img_normalized ** 2)
        features['energy'] = energy / (self.gray_image.shape[0] * self.gray_image.shape[1])

        # 计算熵并归一化
        histogram = np.histogram(self.gray_image, bins=256, range=(0, 256))[0]
        probability = histogram / np.sum(histogram)
        entropy = -np.sum(probability * np.log2(probability + 1e-10))
        features['entropy'] = entropy / np.log2(256)

        # 计算对比度并归一化
        local_variance = np.std(img_normalized) ** 2
        features['contrast'] = local_variance

        return features


def save_results(processor, output_dir='output/'):
    """保存处理结果"""
    import os
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

        # 保存原始Sobel滤波结果
    sobel_result = processor.sobel_filter()
    Image.fromarray(sobel_result).save(f'{output_dir}sobel_result.jpg')

    # 保存原始自定义卷积核滤波结果
    custom_kernel = np.array([[1, 0, -1],
                              [2, 0, -2],
                              [1, 0, -1]])
    custom_result = processor.custom_filter(custom_kernel)
    Image.fromarray(custom_result).save(f'{output_dir}custom_filter_result.jpg')

    # 保存高斯平滑后的边缘检测结果
    gaussian_sobel, gaussian_custom = processor.edge_detection_after_gaussian()
    Image.fromarray(gaussian_sobel).save(f'{output_dir}gaussian_sobel_edge.jpg')
    Image.fromarray(gaussian_custom).save(f'{output_dir}gaussian_custom_edge.jpg')

    # 绘制并保存颜色直方图
    hist = processor.color_histogram()
    plt.figure(figsize=(10, 6))
    colors = ['red', 'green', 'blue']
    for i, color in enumerate(colors):
        plt.plot(hist[:, i], color=color, label=color.upper())
    plt.title('Color Histogram')
    plt.xlabel('Pixel Intensity')
    plt.ylabel('Frequency')
    plt.legend()
    plt.savefig(f'{output_dir}color_histogram.png')
    plt.close()

    # 保存纹理特征
    texture_features = processor.extract_texture_features()
    np.save(f'{output_dir}texture_features.npy', texture_features)

    # 打印归一化后的纹理特征
    print("\n归一化后的纹理特征结果：")
    for feature_name, value in texture_features.items():
        print(f"{feature_name}: {value:.4f}")

    return texture_features


if __name__ == "__main__":
    # 设置输入图像路径
    input_image_path = "test.jpg"

    # 创建图像处理器实例
    processor = ImageProcessor(input_image_path)

    # 处理并保存结果
    texture_features = save_results(processor)
